#include <stdio.h>
#include <stdlib.h>
#include "rbtree.h"

Node *NIL;

// 创建新节点
Node* createNode(int data) {
    Node *node = (Node*)malloc(sizeof(Node));
    node->data = data;
    node->color = RED;  // 新节点默认为红色
    node->left = NIL;
    node->right = NIL;
    node->parent = NIL;
    return node;
}

// 左旋转
void leftRotate(Node **root, Node *x) {
    Node *y = x->right;
    x->right = y->left;
    
    if (y->left != NIL)
        y->left->parent = x;
    
    y->parent = x->parent;
    
    if (x->parent == NIL)
        *root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;
    
    y->left = x;
    x->parent = y;
}

// 右旋转
void rightRotate(Node **root, Node *y) {
    Node *x = y->left;
    y->left = x->right;
    
    if (x->right != NIL)
        x->right->parent = y;
    
    x->parent = y->parent;
    
    if (y->parent == NIL)
        *root = x;
    else if (y == y->parent->right)
        y->parent->right = x;
    else
        y->parent->left = x;
    
    x->right = y;
    y->parent = x;
}

// 红黑树删除后的修复函数
void deleteFixup(Node **root, Node *x) {
    while (x != *root && x->color == BLACK) {
        if (x == x->parent->left) {
            Node *w = x->parent->right;  // 兄弟节点
            
            // 情况1：兄弟节点是红色
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                leftRotate(root, x->parent);
                w = x->parent->right;
            }
            
            // 情况2：兄弟节点是黑色，且其两个子节点都是黑色
            if (w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                // 情况3：兄弟节点是黑色，左子节点是红色，右子节点是黑色
                if (w->right->color == BLACK) {
                    w->left->color = BLACK;
                    w->color = RED;
                    rightRotate(root, w);
                    w = x->parent->right;
                }
                
                // 情况4：兄弟节点是黑色，右子节点是红色
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->right->color = BLACK;
                leftRotate(root, x->parent);
                x = *root;  // 退出循环
            }
        } else {
            // 与上面的情况对称，只是将left和right交换，leftRotate和rightRotate交换
            Node *w = x->parent->left;
            
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                rightRotate(root, x->parent);
                w = x->parent->left;
            }
            
            if (w->right->color == BLACK && w->left->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                if (w->left->color == BLACK) {
                    w->right->color = BLACK;
                    w->color = RED;
                    leftRotate(root, w);
                    w = x->parent->left;
                }
                
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->left->color = BLACK;
                rightRotate(root, x->parent);
                x = *root;
            }
        }
    }
    x->color = BLACK;
}

// 查找最小值节点
Node* minimum(Node *node) {
    while (node->left != NIL)
        node = node->left;
    return node;
}

// 替换子节点
void transplant(Node **root, Node *u, Node *v) {
    if (u->parent == NIL)
        *root = v;
    else if (u == u->parent->left)
        u->parent->left = v;
    else
        u->parent->right = v;
    v->parent = u->parent;
}

// 红黑树删除节点
void deleteNode(Node **root, Node *z) {
    Node *y = z;//最终指向要被实际删除的节点
    Node *x;//指向y被删除后，替代y位置的子节点
    int yOriginalColor = y->color;
    
    // 情况1：z只有一个子节点或没有子节点
    if (z->left == NIL) {
        x = z->right;
        transplant(root, z, z->right);
    } else if (z->right == NIL) {
        x = z->left;
        transplant(root, z, z->left);
    } else {
        // 情况2：z有两个子节点，找到后继节点
        y = minimum(z->right);
        yOriginalColor = y->color;
        x = y->right;// x是y被删除后，替代y位置的节点（y最多只有右孩子）
        
        // 处理后继节点与z的关系
        if (y->parent == z)//NIL节点场景的防御性处理
            x->parent = y;
        else {
            transplant(root, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        
        transplant(root, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    
    free(z);
    
    // 如果删除的是黑色节点，可能破坏红黑树性质，需要修复
    if (yOriginalColor == BLACK)
        deleteFixup(root, x);
}

// 查找节点
Node* search(Node *root, int data) {
    if (root == NIL || root->data == data)
        return root;
    
    if (data < root->data)
        return search(root->left, data);
    else
        return search(root->right, data);
}

// 插入修复函数（为了完整实现红黑树，需要插入功能）
void insertFixup(Node **root, Node *z) {
    while (z->parent->color == RED) {
        if (z->parent == z->parent->parent->left) {
            Node *y = z->parent->parent->right;  // 叔叔节点
            
            // 情况1：叔叔是红色
            if (y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent->color = RED;
                z = z->parent->parent;
            } else {
                // 情况2：叔叔是黑色，且z是右孩子
                if (z == z->parent->right) {
                    z = z->parent;
                    leftRotate(root, z);
                }
                
                // 情况3：叔叔是黑色，且z是左孩子
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                rightRotate(root, z->parent->parent);
            }
        } else {
            // 与上面的情况对称
            Node *y = z->parent->parent->left;
            
            if (y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent->color = RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    rightRotate(root, z);
                }
                
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                leftRotate(root, z->parent->parent);
            }
        }
    }
    (*root)->color = BLACK;
}

// 插入节点
void insertNode(Node **root, int data) {
    Node *z = createNode(data);
    Node *y = NIL;//作父结点用
    Node *x = *root;//利用x找到叶子结点，此时父结点没有左孩或右孩
    
    // 找到插入位置
    while (x != NIL) {
        y = x;
        if (z->data < x->data)
            x = x->left;
        else
            x = x->right;
    }
    
    z->parent = y;
    if (y == NIL)
        *root = z;
    else if (z->data < y->data)
        y->left = z;
    else
        y->right = z;
    
    // 插入后修复红黑树性质
    insertFixup(root, z);
}

// 中序遍历
void inorder(Node *root) {
    if (root != NIL) {
        inorder(root->left);
        printf("%d(%s) ", root->data, root->color == RED ? "R" : "B");
        inorder(root->right);
    }
}

// 初始化NIL节点
void initializeNIL() {
    NIL = (Node*)malloc(sizeof(Node));
    NIL->color = BLACK;
    NIL->left = NIL;
    NIL->right = NIL;
    NIL->parent = NIL;
}